import os
import json
from tqdm import tqdm
from .writer_labels_embeddings import LabelEmbeddingWriter

def process_label_embeddings_nhanes(
    labels_json_path: str,
    output_dir: str,
    model_name: str,
    tensor_parallel_size: int,
    batch_size: int = 1024,
):
    """
    Processes embeddings for labels from JSON files and saves them.
    """
    print("Loading labels from JSON files...")
    label_items = []
    for filename in os.listdir(labels_json_path):
        if filename.endswith(".json"):
            key = filename[:-5]
            file_path = os.path.join(labels_json_path, filename)
            with open(file_path, 'r', encoding='utf-8') as f:
                label_list = json.load(f)
            label_items.append((key, label_list))
    print(f"Loaded {len(label_items)} label sets.")

    print("Writing embeddings for labels...")
    with LabelEmbeddingWriter(
        output_dir,
        model_name,
        tensor_parallel_size=tensor_parallel_size
    ) as embedding_writer:
        for i in range(0, len(label_items), batch_size):
            progress = (i + batch_size) / len(label_items)
            print(f"Progress: {progress:.2%}")
            batch = label_items[i:i + batch_size]
            embedding_writer.embedding_and_write_batch(batch)


    print("Finished processing all labels.")